import torch
from torch.utils import benchmark

cases = [
    (torch.float64, 1),
    (torch.float32, 50),
    (torch.float16, 100),
    # (torch.int16, 100),
    (torch.int8, 100),
]
for typ, repeat in cases:
    n = 1024 * 16
    a = torch.randn(n, n).type(typ).cuda()
    b = torch.randn(n, n).type(typ).cuda()

    t = benchmark.Timer(
        stmt='a @ b',
        globals={'a': a, 'b': b})

    x = t.timeit(repeat)
    print(typ, 2*n**3 / x.median /1e12, 'tflops')
